# -*- coding: utf-8 -*-
# @Time    : 2025/5/21 上午9:40
# @Author  : sjh
# @Site    : 
# @File    : Filter.py
# @Comment :
import time

import cv2
import numpy as np
cv2.setNumThreads(4)
from scipy.ndimage import distance_transform_edt
from scipy.spatial import cKDTree
print(cv2.__version__)
def fill_depth_inpaint(depth_frame, left=64, right=0, bottom=0):
    # 1. 创建一个初始掩膜，值为0的像素为True，否则为False
    mask0 = (depth_frame <= 0)

    height, width = depth_frame.shape
    mask1 = np.zeros_like(depth_frame, dtype=bool)
    mask1[0:height-bottom, left:width-right] = True
    mask = mask0 & mask1
    mask = np.uint8(mask * 255)
    # 使用 cv2.inpaint() 进行修复
    inpaint_radius = 3 # 修复半径，可以根据需要调整
    depth_frame = cv2.inpaint(depth_frame, mask, inpaint_radius, cv2.INPAINT_TELEA)

    return depth_frame

def fill_depth_inpaint_ROI(depth_frame, left=64, right=0, bottom=0):
    # 深度图尺寸
    height, width = depth_frame.shape

    # 定义 ROI 区域（感兴趣区域）
    roi_top = 0
    roi_bottom = height - bottom
    roi_left = left
    roi_right = width - right

    # 截取 ROI 区域
    roi = depth_frame[roi_top:roi_bottom, roi_left:roi_right].copy()

    # 掩码：0 为需要填补的区域
    mask = (roi <= 0).astype(np.uint8) * 255
    print("待修复像素数：", np.sum(mask > 0))
    # 若无0值，直接返回原图
    if not np.any(mask):
        return depth_frame

    roi_filled = cv2.inpaint(roi, mask, 3, cv2.INPAINT_TELEA)

    # 将修复后的 ROI 放回原图
    filled_frame = depth_frame.copy()
    filled_frame[roi_top:roi_bottom, roi_left:roi_right] = roi_filled

    return filled_frame

def fill_depth_nearest(depth_frame, left=200, right=0, bottom=0):
    # 创建一个掩码，表示需要修复的位置（为0 且在感兴趣区域内）
    mask0 = (depth_frame == 0)

    height, width = depth_frame.shape
    mask1 = np.zeros_like(depth_frame, dtype=bool)
    mask1[0:height-bottom, left:width-right] = True

    mask = mask0 & mask1  # 最终掩码：只修复感兴趣区域内的0值
    filled = depth_frame.copy()

    # 最近点填充
    distance, indices = distance_transform_edt(mask, return_indices=True)
    filled[mask] = depth_frame[tuple(indices[:, mask])]
    return filled

def fill_depth_nearest_ROI(depth_frame, left=64, right=0, bottom=0):
    # 深度图尺寸
    height, width = depth_frame.shape

    # 定义 ROI 区域（感兴趣区域）
    roi_top = 0
    roi_bottom = height - bottom
    roi_left = left
    roi_right = width - right

    # 截取 ROI 区域
    roi = depth_frame[roi_top:roi_bottom, roi_left:roi_right].copy()

    # 掩码：0 为需要填补的区域
    mask = roi <= 0
    print("待修复像素数：", np.sum(mask > 0))

    # 若无0值，直接返回原图
    if not np.any(mask):
        return depth_frame
    # 最近邻填充
    distance, indices = distance_transform_edt(mask, return_indices=True)
    roi_filled = roi.copy()
    roi_filled[mask] = roi[tuple(indices[:, mask])]

    # 将修复后的 ROI 放回原图
    filled_frame = depth_frame.copy()
    filled_frame[roi_top:roi_bottom, roi_left:roi_right] = roi_filled

    return filled_frame

def fast_kdtree_fill(depth):
    yx = np.argwhere(depth != 0)
    values = depth[depth != 0]
    tree = cKDTree(yx)

    yx_nan = np.argwhere(depth == 0)
    _, idx = tree.query(yx_nan)
    filled = depth.copy()
    filled[depth == 0] = values[idx]
    return filled

def fill_disp_near_keypoints(disp, keypoints, window_size=10, method='nearest'):
    """
    只对每个关键点周围的区域进行修复，可选填充方式
    :param disp: 输入视差图 (2D numpy array)
    :param keypoints: 关键点列表，每个元素为 (x, y)
    :param window_size: 修复窗口大小（必须为奇数）
    :param method: 填充方法 'nearest' | 'mean' | 'median' | 'inpaint'
    :return: 修复后的视差图
    """
    # assert window_size % 2 == 1, "window_size 必须为奇数"
    assert method in ['nearest', 'mean', 'median', 'inpaint'], "未知填充方法"

    half_w = window_size // 2
    filled = disp.copy()

    def fill_patch_nearest(patch, mask):
        if np.any(mask):
            _, indices = distance_transform_edt(mask, return_indices=True)
            patch[mask] = patch[tuple(indices[:, mask])]
        return patch

    def fill_patch_mean(patch, mask):
        valid = patch[~mask]
        if valid.size > 0:
            patch[mask] = np.mean(valid)
        return patch

    def fill_patch_median(patch, mask):
        valid = patch[~mask]
        if valid.size > 0:
            patch[mask] = np.median(valid)
        return patch

    def fill_patch_inpaint(patch, mask):
        if np.any(mask):
            mask_uint8 = mask.astype(np.uint8)
            patch = cv2.inpaint(patch.astype(np.float32), mask_uint8, 3, cv2.INPAINT_TELEA)
        return patch

    ref_disp = disp[keypoints[19][1], keypoints[19][0]]
    max_diff = 5
    for pt in keypoints:
        # 如果视差图的值为大于0，则不进行修复
        if disp[pt[1], pt[0]] > 0:
            continue
        x, y = int(pt[0]), int(pt[1])
        x1 = max(0, x - half_w)
        x2 = min(disp.shape[1], x + half_w + 1)
        y1 = max(0, y - half_w)
        y2 = min(disp.shape[0], y + half_w + 1)

        patch = filled[y1:y2, x1:x2]
        mask = (np.abs(patch - ref_disp) > max_diff) | (patch <= 0)


        # mask = (patch <= 0)

        if np.any(mask):
            if method == 'nearest':
                patch = fill_patch_nearest(patch, mask)
            elif method == 'mean':
                patch = fill_patch_mean(patch, mask)
            elif method == 'median':
                patch = fill_patch_median(patch, mask)
            elif method == 'inpaint':
                patch = fill_patch_inpaint(patch, mask)

            filled[y1:y2, x1:x2] = patch

    return filled


def show_disp(disp):
    norm = ((disp - disp.min()) / (disp.max() - disp.min()) * 255).astype(np.uint8)
    depth_colormap = cv2.applyColorMap(norm, cv2.COLORMAP_PLASMA)
    return depth_colormap
if __name__ == "__main__":
    import timeit
    # 读取深度图（16位）
    depth_map = cv2.imread('/home/orangepi/sjh/SGBM_port_sjh_158/data/disp.png', cv2.IMREAD_UNCHANGED)
    print(np.min(depth_map), np.max(depth_map), depth_map.dtype)
    print(depth_map.shape)

    t1 = time.time()
    for i in range(10):
        inpainted_depth = fill_depth_nearest_ROI(depth_map)
    print("Time taken for inpainting:", time.time()-t1)
    print(timeit.timeit(lambda: fill_depth_inpaint(depth_map), number=10))
    print(timeit.timeit(lambda: fill_depth_inpaint_ROI(depth_map), number=10))
    print(timeit.timeit(lambda: fill_depth_nearest(depth_map), number=10))
    print(timeit.timeit(lambda: fill_depth_nearest_ROI(depth_map), number=10))
    print(timeit.timeit(lambda: fast_kdtree_fill(depth_map), number=10))



    # 归一化到 0~255，用于可视化（避免16位深度显示不正确）
    depth_norm = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX)
    depth_norm = depth_norm.astype(np.uint8)

    inpainted_norm = cv2.normalize(inpainted_depth, None, 0, 255, cv2.NORM_MINMAX)
    inpainted_norm = inpainted_norm.astype(np.uint8)


    depth_norm = show_disp(depth_norm)
    inpainted_norm = show_disp(inpainted_norm)
    # 拼接显示
    combined_img = np.hstack((depth_norm, inpainted_norm))
    # cv2.imshow('Original vs Inpainted Depth Map (Normalized)', combined_img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()


